import matplotlib.pyplot as plt
import numpy as np
from xray_scan import XRayScan
from xray_scan_collection import XRayScanCollection
from xray_scan_collection_fft import XRayScanCollectionFFT

class XRayScanPlotter:
    def __init__(self, xray_scan_collection_fft):
        """
        初始化XRayScanPlotter对象。

        :param xray_scan_collection_fft: XRayScanCollectionFFT对象
        """
        self.xray_scan_collection_fft = xray_scan_collection_fft

    def plot_polar_magnitude_spectrum(self):
        """
        在极坐标系中绘制XRayScanCollectionFFT中的各个XRayScan对象的幅度谱。
        """
        plt.figure(figsize=(8, 8))
        ax = plt.subplot(111, projection='polar')  # 创建极坐标子图

        for scan in self.xray_scan_collection_fft.xray_scan_collection.scans:
            # 获取角度并转换为弧度
            angle_rad = np.radians(scan.angle_deg)
            n = len(scan.sensor_data)
            r_values = np.arange(-scan.origin_index, n - scan.origin_index)

            # 计算幅度谱
            magnitude_spectrum = [np.abs(complex_val) for complex_val in scan.sensor_data]

            # 将极坐标数据绘制
            theta = np.full_like(r_values, angle_rad)  # 创建角度数组
            ax.plot(theta, magnitude_spectrum, marker='o', label=f"Angle {scan.angle_deg}°")

        ax.set_title("Magnitude Spectrum in Polar Coordinates")
        ax.set_rlabel_position(-22.5)  # 调整径向标签位置
        ax.legend()
        plt.show()

def main():
    # 创建几个XRayScan对象
    scan1 = XRayScan(30, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan2 = XRayScan(60, [0, 1, 0, 3, 0, 2, 0, 1, 0])
    scan3 = XRayScan(90, [0, 1, 0, 3, 0, 2, 0, 1, 0])

    # 创建XRayScanCollection对象并添加扫描
    collection = XRayScanCollection()
    collection.add_scan(scan1)
    collection.add_scan(scan2)
    collection.add_scan(scan3)

    # 创建XRayScanCollectionFFT对象，并计算傅里叶变换
    fft_collection = XRayScanCollectionFFT(collection)
    fft_collection.compute_fft()

    # 创建XRayScanPlotter对象并绘制图像
    plotter = XRayScanPlotter(fft_collection)
    plotter.plot_polar_magnitude_spectrum()

if __name__ == "__main__":
    main()